package com.db.router.dynamic;

import com.db.router.context.ConnectionContext;
import com.db.router.context.DBContextHolder;
import com.db.router.context.TransactionContext;
import org.apache.commons.lang.StringUtils;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * @author tianjiashu（mac）
 * @ClassName DynamicDataSource.java
 * @Description 重写获取连接方法，获取连接后将连接存起来。
 * @createTime 2023年07月12日 22:15:00
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    public Connection getConnection() throws SQLException {
        return getConnection("","");
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        DataSource dataSource = super.determineTargetDataSource();
        CustomConnection connection = ConnectionContext.getConnection(dataSource);
        if(connection != null){
            // 复用连接
            return connection;
        }

        if(StringUtils.isNotBlank(username) && StringUtils.isNotBlank(password)){
            connection = new CustomConnection(dataSource.getConnection(username, password));
        }else {
            connection = new CustomConnection(dataSource.getConnection());
        }
        //如果事务标志开启，就禁止自动提交事务，并保存连接
        if (TransactionContext.isOpenTran()) {
            connection.setAutoCommit(false);
            ConnectionContext.setConnection(dataSource,connection);
        }
        return connection;
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return "db" + DBContextHolder.getDBKey();
    }

}
